package com.shiyi.shiyiaiagent.advisor;

import org.springframework.ai.chat.client.advisor.api.*;
import reactor.core.publisher.Flux;

import java.util.HashMap;

/**
 * 自定义 Re 2 Advisor
 * 可提高大型语言模型的推理能力
 */
public class ReReadingAdvisor implements CallAroundAdvisor, StreamAroundAdvisor {

  private AdvisedRequest before(AdvisedRequest advisedRequest) {
    HashMap<String, Object> advisorUserParams = new HashMap<>(advisedRequest.userParams());
    advisorUserParams.put("re2_input_query", advisedRequest.userText());

    return AdvisedRequest.from(advisedRequest)
      .userText("""
        {re2_input_query}
        Read the question again: {re2_input_query}
        """)
      .userParams(advisorUserParams)
      .build();
  }
  @Override
  public AdvisedResponse aroundCall(AdvisedRequest advisedRequest, CallAroundAdvisorChain chain) {

    return  chain.nextAroundCall(this.before(advisedRequest));
  }

  @Override
  public Flux<AdvisedResponse> aroundStream(AdvisedRequest advisedRequest, StreamAroundAdvisorChain chain) {
    return chain.nextAroundStream(this.before(advisedRequest));
  }

  @Override
  public String getName() {
    return this.getClass().getSimpleName();
  }

  @Override
  public int getOrder() {
    return 0;
  }
}
